package com.whosly.mysql.protocol.packet.impl;

import com.whosly.mysql.protocol.Capability;
import com.whosly.mysql.protocol.config.MySQLProtocolConfig;
import com.whosly.mysql.protocol.jdbc.MySQLMessage;
import com.whosly.mysql.protocol.packet.jdbc.IAuthPacket;
import com.whosly.mysql.protocol.packet.jdbc.MySQLPacket;
import com.whosly.mysql.protocol.packet.util.NettyBufferUtil;
import com.whosly.mysql.protocol.util.ByteUtils;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import org.apache.commons.lang3.StringUtils;

import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import static com.whosly.mysql.protocol.Capability.ConnectAttrs;
import static com.whosly.mysql.protocol.Capability.ConnectWithDatabase;
import static com.whosly.mysql.protocol.Capability.PluginAuth;

/**
 *
 * @author 风杨
 * @see <a href="yueny09@163.com">mailTo: yueny09@163.com</a>
 */
public class AuthPacket extends MySQLPacket implements IAuthPacket,
        Serializable {
    private static final long serialVersionUID = -3968420616392862050L;

    private static final byte[] FILLER23 = new byte[23];

    // capabilities
    public int clientFlags;
    // highThreeBytes
    public int maxPacketSize;
    public byte charsetIndex;
    public byte[] extra;// from FILLER(23)
    public String user;

    private byte[] encryptedPassword;

    public String database;

    // 如 mysql_native_password
    public String clientAuthPlugin;

    /**
     * 客户端属性
     *
     * 值不区分大小写。 key固定
     *
     * connectionAttributes=_source:[test/dms],_client_auth_type:[password/token]";
     */
    public Map<String, String> connectAttrs;

    @Override
    public void read(byte[] bodyBytes)
    {
        MySQLMessage mm = new MySQLMessage(bodyBytes);
        int position = mm.position();
        clientFlags = mm.readInt();
        if (!Capability.isEnabled(clientFlags, Capability.Speaks4_1Protocl_NewFlag)) {
            // do nothing
        } else {
            mm.position(position);

            clientFlags = mm.readInt();
            maxPacketSize = mm.readInt();
            charsetIndex = mm.read();

            int current = mm.position();
            int len = (int) mm.readLength();
            if (len > 0 && len < FILLER23.length) {
                byte[] ab = new byte[len];
                System.arraycopy(mm.bytes(), mm.position(), ab, 0, len);
                this.extra = ab;
            }
            // 跳过23个填充字节
            mm.position(current + 23);

            user = mm.readStringWithNull();

            if (Capability.isEnabled(this.clientFlags, Capability.PluginAuthLenEncClientData)
                    || Capability.isEnabled(this.clientFlags, Capability.CanDo4_1Authentication)) {
                setEncryptedPassword(mm.readBytesWithLength());
            } else {
                setEncryptedPassword(mm.readBytesWithNull());
            }

            // 读取连接的数据库名称
            if (Capability.isEnabled(clientFlags, ConnectWithDatabase) && mm.hasRemaining()) {
                database = mm.readStringWithNull();
            }

            // 读取authentication plugin的名字
            if (Capability.isEnabled(clientFlags, PluginAuth) && mm.hasRemaining()) {
                clientAuthPlugin = mm.readStringWithNull();
            }

            // 读取客户端的一些属性
            if (Capability.isEnabled(clientFlags, ConnectAttrs) && mm.hasRemaining()) {
                connectAttrs = new HashMap<>(8);
                long totalLength = mm.readLength();
                long readLength = 0;
                while (readLength < totalLength) {
                    int startPos = mm.position();
                    String key = mm.readStringWithLength();
                    String value = mm.readStringWithLength();
                    int endPos = mm.position();
                    connectAttrs.put(key, value);
                    readLength += (endPos - startPos);
                }
            }
        }
    }

    public void write(Channel c) {
        // default init 256,so it can avoid buff extract
        ByteBuf buffer = c.alloc().buffer();
        NettyBufferUtil.writeUB3(buffer, calcPacketSize());
        buffer.writeByte(getPackId());

        writeByteBuf(buffer);

        c.writeAndFlush(buffer);
    }

    public void write(ChannelHandlerContext ctx) {
        // default init 256,so it can avoid buff extract
        ByteBuf buffer = ctx.alloc().buffer();
        NettyBufferUtil.writeUB3(buffer, calcPacketSize());
        buffer.writeByte(getPackId());

        writeByteBuf(buffer);

        ctx.writeAndFlush(buffer);
    }

    @Override
    public ByteBuf writeBuf(ByteBuf buffer)
    {
        NettyBufferUtil.writeUB3(buffer, buffer.capacity());
        buffer.writeByte(getPackId());

        writeByteBuf(buffer);

        return buffer;
    }

    @Override
    public ByteBuf writeBuf()
    {
        // 在堆外分配一个ByteBuf，并指定初始容量
        ByteBuf buffer = Unpooled.directBuffer(calcPacketSize());

        return writeBuf(buffer);
    }

    public void writeByteBuf(ByteBuf buffer) {
        NettyBufferUtil.writeUB4(buffer, clientFlags);
        NettyBufferUtil.writeUB4(buffer, maxPacketSize);
        buffer.writeByte(charsetIndex);
        buffer.writeBytes(FILLER23);

        if (user == null) {
            buffer.writeByte((byte) 0);
        } else {
            byte[] userData = user.getBytes();
            NettyBufferUtil.writeWithNull(buffer, userData);
        }

        if (Capability.isEnabled(this.clientFlags, Capability.PluginAuthLenEncClientData)
            || Capability.isEnabled(this.clientFlags, Capability.CanDo4_1Authentication)) {
            if (getEncryptedPassword() == null) {
                buffer.writeByte((byte) 0);
            } else {
                NettyBufferUtil.writeWithLength(buffer, getEncryptedPassword());
            }
        } else {
            if (getEncryptedPassword() == null) {
                buffer.writeByte((byte) 0);
            } else {
                NettyBufferUtil.writeWithNull(buffer, getEncryptedPassword());
            }
        }

        if (Capability.isEnabled(clientFlags, ConnectWithDatabase)) {
            if (database == null) {
                buffer.writeByte((byte) 0);
            } else {
                byte[] databaseData = ByteUtils.getBytes(database);
                NettyBufferUtil.writeWithNull(buffer, databaseData);
            }
        }

        if (Capability.isEnabled(clientFlags, PluginAuth)) {
            if (clientAuthPlugin == null) {
                buffer.writeByte((byte) 0);
            } else {
                byte[] authPlugin = ByteUtils.getBytes(clientAuthPlugin, MySQLProtocolConfig.DEFAULT_CHARSET);
                NettyBufferUtil.writeWithNull(buffer, authPlugin);
            }
        }

        if (Capability.isEnabled(clientFlags, ConnectAttrs)) {
            int totalLength = getConnectAttrsLength();
            if (totalLength > 0) {
                NettyBufferUtil.writeLength(buffer, totalLength);
                for (Map.Entry<String, String> kv : connectAttrs.entrySet()) {
                    if (kv.getKey() == null) {
                        buffer.writeByte((byte) 0);
                    } else {
                        NettyBufferUtil.writeWithLength(buffer, kv.getKey().getBytes());
                    }
                    if (kv.getValue() == null) {
                        buffer.writeByte((byte) 0);
                    } else {
                        NettyBufferUtil.writeWithLength(buffer, kv.getValue().getBytes());
                    }
                }
            }
        }

    }

    @Override
    public int calcPacketSize() {
        int size = 32;// 4+4+1+23;
        size += (user == null) ? 1 : user.length() + 1;
        if (Capability.isEnabled(this.clientFlags, Capability.PluginAuthLenEncClientData)
                || Capability.isEnabled(this.clientFlags, Capability.CanDo4_1Authentication)) {
            size += (getEncryptedPassword() == null) ? 1 : NettyBufferUtil.getLengthVarSize(getEncryptedPassword());
        } else {
            size += (getEncryptedPassword() == null) ? 1 : getEncryptedPassword().length + 1;
        }
        if (Capability.isEnabled(clientFlags, ConnectWithDatabase)) {
            size += (database == null) ? 1 : database.length() + 1;
        }
        if (Capability.isEnabled(clientFlags, PluginAuth)) {
            size += (clientAuthPlugin == null) ? 1 : clientAuthPlugin.length() + 1;
        }
        int totalLength = getConnectAttrsLength();
        if (Capability.isEnabled(clientFlags, ConnectAttrs)) {
            size += NettyBufferUtil.getLengthVarPrefixSize(totalLength);
            size += totalLength;
        }
        return size;
    }

    /**
     * 获取kv对的总长度
     * @return
     */
    private int getConnectAttrsLength() {
        int totalLength = 0;
        if (connectAttrs == null || connectAttrs.isEmpty()) {
            return totalLength;
        }
        for (Map.Entry<String, String> kv : connectAttrs.entrySet()) {
            if (kv.getKey() == null) {
                totalLength += 1;
            } else {
                totalLength += NettyBufferUtil.getLengthVarSize(kv.getKey().getBytes());
            }
            if (kv.getValue() == null) {
                totalLength += 1;
            } else {
                totalLength += NettyBufferUtil.getLengthVarSize(kv.getValue().getBytes());
            }
        }
        return totalLength;
    }

    @Override
    public byte getCharset() {
        return charsetIndex;
    }

    @Override
    public Optional<String> getUsername() {
        return StringUtils.isEmpty(user) ? Optional.empty() : Optional.ofNullable(user);
    }

    public void setEncryptedPassword(byte[] encryptedPassword) {
        this.encryptedPassword = encryptedPassword;
    }

    @Override
    public byte[] getEncryptedPassword() {
        return encryptedPassword;
    }

    @Override
    public Optional<String> getDatabase() {
        return StringUtils.isEmpty(database) ? Optional.empty() : Optional.ofNullable(database);
    }

    @Override
    public byte getPacketId() {
        return getPackId();
    }

    @Override
    public String getPacketInfo() {
        return "MySQL Authentication Packet";
    }
}
